from typing import List
from fastapi import FastAPI, WebSocket, WebSocketException
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
import json
import asyncio
import os
from typing import Dict
import uvicorn
from fastapi import Request
from zhipuai import ZhipuAI

app = FastAPI()
tasks = []
task_lock = asyncio.Lock()
active_connections: List[WebSocket] = []

zhipu_api_key = ""
zhipu_api_file = "zhipuapi.txt"
if os.path.exists(zhipu_api_file):
    try:
        with open(zhipu_api_file, "r", encoding="utf-8") as f:
            zhipu_api_key = f.read().strip()
    except Exception as e:
        print(f"读取zhipuapi.txt失败: {e}")
else:
    print("未找到zhipuapi.txt文件")

# 用户任务字典
user_tasks: Dict[str, list] = {}
user_tasks_file = "user_tasks.json"

# 加载本地文件中的用户任务数据
if os.path.exists(user_tasks_file):
    try:
        with open(user_tasks_file, "r", encoding="utf-8") as f:
            user_tasks = json.load(f)
    except Exception as e:
        print(f"加载用户任务文件失败: {e}")

app.mount("/static", StaticFiles(directory="static"), name="static")

@app.post("/tasks")
async def add_task(request: Request):
    try:
        json_data = await request.json()
        user_name = json_data.get("UserName")
        task_list = json_data.get("TaskLists")
        if not user_name or not isinstance(task_list, list):
            return {"error": "缺少UserName或TaskLists格式不正确"}

        async with task_lock:
            user_tasks[user_name] = task_list
            try:
                with open(user_tasks_file, "w", encoding="utf-8") as f:
                    json.dump(user_tasks, f, ensure_ascii=False, indent=2)
            except Exception as e:
                print(f"保存用户任务文件失败: {e}")
            for ws in active_connections:
                await ws.send_text(json.dumps({"user_tasks": user_tasks}))
        return {"message": "接收到了新的任务列表"}
    except Exception as e:
        return {"error": f"Json格式错误: {e}"}

@app.websocket("/ws_tasks_list")
async def websocket_tasks_list(websocket: WebSocket):
    await websocket.accept()
    active_connections.append(websocket)
    try:
        async with task_lock:
            await websocket.send_text(json.dumps({"user_tasks": user_tasks}))
        while True:
            await websocket.receive_text()
    except Exception as e:
        print(f"Error: {e}")
    finally:
        active_connections.remove(websocket)
        await websocket.close()

@app.post("/api/submit")
async def api_submit(request: Request):
    try:
        json_data = await request.json()
        mission_content = json_data.get("data")
        print(f"接收到的任务内容: {mission_content}")
        client = ZhipuAI(api_key=zhipu_api_key)
        response = client.chat.completions.create(
            model="glm-4-plus",
            messages=[
                {"role": "system", "content": "你是一个辅助用户进行计划和任务管理的智能助手，用户将传递给你一个Json格式的任务列表，你需要根据这些任务生成一个智能规划建议。你生成的建议应当最好控制在100字以内，包含用户首先应当完成的、优先级较高的任务，和用户应当注意的长期任务，以及提供给用户完成任务的一些具体方式建议。不要包含Markdown语法，例如星号强调，只要给出纯文本字符串即可。注意被标注为已完成/废弃的任务，这代表用户并不关注它们，无需针对它们生成建议。"},
                {"role": "user", "content": mission_content}
            ],
        )
        if response and response.choices:
            print(f"生成的建议: {response.choices[0].message.content}")
            return {"suggestion": response.choices[0].message.content}
        else:
            return {"suggestion": "未能生成建议"}
    except Exception as e:
        return {"error": f"提交失败: {e}"}

@app.get("/", response_class=HTMLResponse)
async def get_root():
    return FileResponse("static/index.html")

if __name__ == "__main__":
    uvicorn.run(app, host="localhost", port=9961)